import uuid
from datetime import datetime, timedelta
from typing import Optional, Annotated, List
import json

import jwt
from fastapi import Depends
from fastapi.security import OAuth2PasswordBearer
from jwt import InvalidTokenError
from pydantic import BaseModel
from sqlalchemy.ext.asyncio import AsyncSession
from starlette.requests import Request
from starlette.responses import JSONResponse

from src.base.context import web_login_user
from src.base.enums import DataScopeType
from src.config.db import redisClient
from src.dao.sys_dept_dao import get_dept_with_childes_by_ids
from src.dao.sys_role_dao import get_role_list_by_user_id

JWT_SECRET = '123456'
JWT_ALGORITHM = 'HS512'
JWT_EXPIRATION = 86400
JWT_TOKEN_PREFIX = 'Bearer'
JWT_CACHE_KEY_PREFIX = 'fast-test:login_user:'
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/auth/login")


class DataScope(BaseModel):
    dept_ids: Optional[List[int]] = None
    user_ids: Optional[List[int]] = None


class LoginUser(BaseModel):
    user_id: Optional[int] = None
    user_name: Optional[str] = None
    nick_name: Optional[str] = None
    tenant_id: Optional[int] = None
    dept_ids: Optional[List[int]] = None
    data_scope: Optional[DataScope] = DataScope()
    last_login_time: Optional[datetime] = None
    expire_time: Optional[float] = None
    user_type: Optional[str] = None
    login_ip: Optional[str] = None
    cache_key: Optional[str] = None


def get_cache_key(cache_key: str) -> str:
    return f'{JWT_CACHE_KEY_PREFIX}{cache_key}'


# 生成token
async def create_access_token(login_user: LoginUser, dbSession: AsyncSession) -> str:
    user_ids, dept_ids = await set_data_scope(login_user.user_id, login_user.dept_ids, dbSession)
    login_user.data_scope.user_ids = user_ids
    login_user.data_scope.dept_ids = dept_ids
    login_user.cache_key = str(uuid.uuid4())
    login_user.expire_time = JWT_EXPIRATION
    redisClient.set(get_cache_key(login_user.cache_key), login_user.model_dump_json(), login_user.expire_time)
    return jwt.encode({'id': login_user.cache_key}, JWT_SECRET, algorithm=JWT_ALGORITHM)


# 校验token
def verify_access_token(token: str) -> LoginUser | JSONResponse:
    try:
        payload = jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
        if not payload or 'id' not in payload or not payload['id']:
            # raise AuthException(name='Auth Error', message='token 不合法')
            return JSONResponse(status_code=401, content={"detail": "身份验证失败，请重新登录！！！"},
                                headers={"WWW-Authenticate": "Bearer"})
        login_user = redisClient.get(get_cache_key(payload['id']))
        if login_user is None:
            # raise AuthException(name='Auth Error', message='token已过期')
            return JSONResponse(status_code=401, content={"detail": "token已过期，请重新登录！！！"},
                                headers={"WWW-Authenticate": "Bearer"})
    except InvalidTokenError:
        # raise AuthException(name='Auth Error', message='token 不合法')
        return JSONResponse(status_code=401, content={"detail": "token已过期，请重新登录！！！"},
                            headers={"WWW-Authenticate": "Bearer"})
    return LoginUser.model_validate(json.loads(login_user))


# 认证依赖
async def auth_depend(request: Request, token: Annotated[str, Depends(oauth2_scheme)]):
    login_user = verify_access_token(token)
    if isinstance(login_user, JSONResponse):
        return login_user
    request.state.login_user = login_user
    web_login_user.set(login_user)


async def set_data_scope(user_id: int, dept_ids: List[int], dbSession: AsyncSession):
    user_ids = list()
    roles = await get_role_list_by_user_id(dbSession, user_id)
    data_scope_list = [role.data_scope for role in roles]
    if DataScopeType.ALL.value not in data_scope_list:
        for role in roles:
            if role.data_scope == DataScopeType.CUSTOM.value:
                dept_ids.extend(role.custom_dept_ids)
                user_ids.extend(role.custom_user_ids)
            elif role.data_scope == DataScopeType.DEPT_AND_CHILD.value:
                dept_list = await get_dept_with_childes_by_ids(dbSession, dept_ids)
                dept_ids.extend([dept.dept_id for dept in dept_list])
            elif role.data_scope == DataScopeType.DEPT_ONLY.value:
                dept_ids.extend(dept_ids)
            elif role.data_scope == DataScopeType.SELF.value:
                user_ids.append(user_id)
    return list(set(user_ids)), list(set(dept_ids))
